package org.west.sky.scripture.ratelimiter;

import java.time.LocalTime;
import java.util.Set;
import java.util.TreeMap;

/**
 * @author chz
 * @date 2022/3/18
 * @description: 滑动日志限流算法
 */
public class SlideLogRateLimiter {

    /**
     * 阈值
     */
    private Integer qps = 2;
    /**
     * 记录请求时间戳和数量
     */
    private TreeMap<Long, Long> treeMap = new TreeMap<>();
    /**
     * 清理记录时间 60s
     */
    private long clearTime = 60 * 1000;

    public SlideLogRateLimiter(int qps) {
        this.qps = qps;
    }

    public synchronized boolean tryAcquire() {
        long now = System.currentTimeMillis();
        //清理过期的数据老数据 最短60s清理一次
        if (!treeMap.isEmpty() && (treeMap.firstKey() - now) > clearTime) {
            Set<Long> keySet = treeMap.subMap(0L, now - 1000).keySet();
            for (Long aLong : keySet) {
                treeMap.remove(aLong);
            }
        }
        //计算当前请求次数
        int sum = 0;
        for (Long value : treeMap.subMap(now - 1000, now).values()) {
            sum += value;
        }
        //超过QPS直接返回
        if (sum + 1 > qps) {
            return false;
        }
        //记录本次请求
        if (treeMap.containsKey(now)) {
            treeMap.compute(now, (k, v) -> v + 1);
        } else {
            treeMap.put(now, 1L);
        }
        return sum <= qps;
    }

    public static void main(String[] args) throws InterruptedException {
        SlideLogRateLimiter logRateLimiter = new SlideLogRateLimiter(3);
        for (int i = 0; i < 20; i++) {
            Thread.sleep(250);
            LocalTime time = LocalTime.now();
            if (logRateLimiter.tryAcquire()) {
                System.out.println(time + " 做点什么");
            }else {
                System.out.println(time + " 被限流了");
            }
        }
    }
}
